import pandas as pd
import numpy as np
import os
from LMdata.LMdataset import idx2kp
from utils.metrix import cal_NEscore

true_csv = pd.read_csv('./val1/true_val.csv')
files = [
    # 'GCN512_bs4_5.25.csv',
    # 'GCN336_bs4_5.85_336pred.csv',
    # 'GCN336_bs4-5.85_2.csv'
]
cates = np.load('./val1/val_cates.npy')

pred_xyv = np.zeros((len(files), true_csv.shape[0], 24, 3), dtype=np.int)
true_xyv = np.zeros((true_csv.shape[0], 24, 3), dtype=np.int)

for j,file in enumerate(files):
    df = pd.read_csv(os.path.join('./val1', file))

    for i in range(24):
        pred_xyv[j, :, i, 0] = df[idx2kp[i]].str.split('_').str[0].astype(int)  # x coord [-1,512]
        pred_xyv[j, :, i, 1] = df[idx2kp[i]].str.split('_').str[1].astype(int)  # y coord [-1,512]
        pred_xyv[j, :, i, 2] = df[idx2kp[i]].str.split('_').str[2].astype(int)  # vis  [-1, 0, 1]


for i in range(24):
    true_xyv[:, i, 0] = true_csv[idx2kp[i]].str.split('_').str[0].astype(int)  # x coord [-1,512]
    true_xyv[:, i, 1] = true_csv[idx2kp[i]].str.split('_').str[1].astype(int)  # y coord [-1,512]
    true_xyv[:, i, 2] = true_csv[idx2kp[i]].str.split('_').str[2].astype(int)  # vis  [-1, 0, 1]



pred_xyv = pred_xyv.mean(0).astype(int)

score, details =  cal_NEscore(pred_xyv[:, :, 0:2], true_xyv, cates)


print('val-NE: %.4f%%' % (score * 100))
for key in details.keys():
    print('NE: %.3f%% supports: %d  %s' %
          (details[key][0] * 100, details[key][1], key))

